/*
 * @(#)WebSocketFramer.java 0.1.0 01/27/2012
 * Copyright(c) CS5321
 * Baylor University, Waco, Texas, U.S.A
 * All Rights Reserved.
 */
package edu.baylor.websocket;

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Random;

/**
 * WebSocketFramer implements the data frame according to server 
 * requirement of RFC6544
 * 
 * @version 1.0.0 01/27/2012
 * @author Yanqing Liu
 */
public class WebSocketFramer {

    /* opcode */
    public static final int CONTINUATION = 0x00;
    public static final int TEXT = 0x01;
    public static final int BINARY = 0x02;
    public static final int CLOSE = 0x08;
    public static final int PING = 0x09;
    public static final int PONG = 0x0A;

    /* Mask code and Shift code*/
    public static final int OPCODEMASK = 0x000f;    //Get length field
    public static final int LENGTHMASK = 0x007f;    //Get length field
    public static final int MASKMASK = 0x0080;      //Get mask field
    public static final int MASKSHIFT = 7;      
    public static final int MASKING_KEY_SIZE = 4;
    public static final int EXTENED_PAYLOAD_SIZE = 8;
    public static final int NO_FRAG = 0x80;     //Extension set to 0;
    public static final int UNSIGNED_INT_SIZE = 32767;
    
    private DataInputStream in;         // wrapper for data I/O
    private OutputStream out;           // data output
    
    private static final int MAX_PAYLOAD_LENGTH = 65535;
    private static byte[] message = new byte[MAX_PAYLOAD_LENGTH];
    private static int index = 0;       //indicate the next byte position
    
    private enum State {
        single, frag
    }
        private State currentState = State.single;
    
    /**
     * Create WebSocketFramer from InputStream and OutputStream
     * @param in InputStream from which get handshake message
     * @param out OutputStream to which send handshake message
     */
    public WebSocketFramer(InputStream in, OutputStream out) {
        this.in = new DataInputStream(in);
        this.out = out;
    }
    
    /**
     * Send data frame according to FRC6544
     * @param IWSMessage that need to send
     * @throws IOException when there is a connection problem
     *      
     */
    public void frameMsg(IWSMessage msg) throws IOException {
      
        if (msg.getMessageType() == IWSMessage.WSMessageType.TEXT) {
            sendFrame(msg.getData(), TEXT);
        } else if (msg.getMessageType() == IWSMessage.WSMessageType.BINARY) {
            sendFrame(msg.getData(), BINARY);
        }
    }

    /**
     * Receive data frame according to FRC6544
     * @return message read
     * @throws IOException when there is a connection problem     
     */
    public WSMessage nextMsg() throws IOException {
        int nextByte;
        int opcode;
        int fin;

        WSMessage msg = new WSMessage();

        try {
            nextByte= in.read(); // read 1 bytes                
        } catch (EOFException e) { 
            return null;
        }
        
        fin = nextByte & NO_FRAG;
        
        if (fin == 0) {
            currentState = State.frag;
        } else{
            //Do nothing
        }
        
        /* According to the opcode, handle different message*/
        // TODO: Need to support other opcode
        opcode = nextByte & OPCODEMASK;
        
        getPayload();
        
        while (currentState != State.single){
            continueProcess(opcode);
        }
        
        switch (opcode) {
        case TEXT:      
            msg.type = IWSMessage.WSMessageType.TEXT; 
            break;
        case BINARY:    
            msg.type = IWSMessage.WSMessageType.BINARY; 
            break;
        case CLOSE:     
            sendClose(); 
            return null;
        default:    
            sendClose(); 
            return null;
        }
        copy2Msg(msg);
        
        return msg;
    }
    
    /**
     * Send close control message according to FRC6544
     * @return message read
     * @throws IOException when there is a connection problem     
     */
    public void sendClose() throws IOException {
        ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
        DataOutputStream dataOut = new DataOutputStream(byteStream);
        
        dataOut.write(NO_FRAG | CLOSE);
        dataOut.write(MASKMASK);
        //dataOut.writeShort(code);
        //dataOut.writeShort(reason);
        dataOut.flush();
        out.write(byteStream.toByteArray());
     }
    
    /*
     * Send data frame according to FRC6544
     * @param data application to send
     * @param type the type of the data fram
     * @throws IOException when there is a connection problem     
     */
    private void sendFrame(byte[] data, int type) throws IOException {
    
        int length = data.length;
        int randomInt;
        int opcode = TEXT;
        boolean enableMask = true;
    
        ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
        DataOutputStream dataOut = new DataOutputStream(byteStream);
    
        Random random = new Random(System.currentTimeMillis());
        randomInt = random.nextInt();
    
        byte[] maskKey = {(byte) (randomInt  >> 24),
                          (byte) (randomInt >> 16),
                          (byte) (randomInt >> 8),
                          (byte) (randomInt)};
        byte mask = 0;
    
        opcode = NO_FRAG | type;
        
        dataOut.write(NO_FRAG | opcode);
    
        if (enableMask) {
            mask = (byte) MASKMASK;
        }
        
        /* encode the length */ 
        if (length < 126) {
            dataOut.writeByte(mask | (byte) length);
        } else if ((length >=126) && (length <=UNSIGNED_INT_SIZE)) {
            dataOut.writeByte(mask | 126);
            dataOut.writeShort((short)length);
        } else {
            //TODO: need support the large message
            throw new IOException("Payload length too large!");
        }
        
        if (enableMask) {
            /* encode mask key */
            dataOut.write(maskKey);
            for(int i = 0; i< length ;  i++) {
                data[i] = (byte)(data[i] ^ maskKey[i % 4]);
            }
        }
        
        dataOut.write(data);
        dataOut.flush();
        out.write(byteStream.toByteArray());
    }

    
    /* 
     * If encounter frag, continue receive frame until get 
     * a fin indication  
     */
    private void continueProcess(int type) throws IOException {
        int nextByte;
        int opcode;
        int fin;
    
        try {
            nextByte= in.read(); // read 1 bytes                
        } catch (EOFException e) { 
                throw new IOException("Exception while reading");
        }
        
        fin = nextByte & NO_FRAG;
        
        if (fin != 0){
            currentState = State.single;
        }
        
        /* According to the opcode, handle different message*/
        // TODO: Need to support other opcode
        opcode = nextByte & OPCODEMASK;
        
        if (type != opcode) {
            throw new IOException("Opcode is different.");
        }
        
        getPayload();;
        
        if (currentState == State.frag){

        }
        
    }
    
    /* Get unmasked payload */
    private void getPayload() throws IOException{
        int length;     //used for calculate payload length
        int nextByte;       
        int mask;
        //TODO: support extended payload:int extenedPayloadSize; 
        byte[] maskKey = new byte[MASKING_KEY_SIZE];
        
        nextByte= in.read();
        length = nextByte & LENGTHMASK;
        mask = (nextByte & MASKMASK) >> MASKSHIFT;      
 
        /* Calculate the true payload size */
        if(length < 126) {
        } else if (length == 126) {   
            length = in.readUnsignedShort();
        } else {
            //TODO: Recalculate the payload length;
            throw new IOException("lengh = 127 not supported");
        }
        
        if (mask == 1) { 
            for (int i = 0; i < 4; i++) {
                maskKey[i] = (byte)in.read();  
            }
        }
        
        byte[] payload = new byte[length];

        in.readFully(payload);

        /* Unmask the payload */ 
        for (int i = 0; i < length ; i++) {
            message[index++] = (byte) (payload[i]^maskKey[i % 4]);
        }
    }

    
    /* Convert from vector to bytes*/
    private void copy2Msg(WSMessage msg){
        byte[] byteArray = new byte[index];
    
        for (int i = 0; i < index; i++){
            byteArray[i] = message[i];
        }
        
        msg.payload = byteArray;
    
        /* Clear message buffer */
        index = 0;
    }
}
